#!/usr/bin/python
# encoding=utf-8
from StringIO import StringIO
import paramiko
import time
class SshClient:
    TIMEOUT = 4

    def __init__(self, host, port, username, password, key=None, passphrase=None):
        self.username = username
        self.password = password
        self.client = paramiko.SSHClient()
        self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        if key is not None:
            key = paramiko.RSAKey.from_private_key(StringIO(key), password=passphrase)
        self.client.connect(host, port, username=username, password=password, pkey=key, timeout=self.TIMEOUT)

    def close(self):
        if self.client is not None:
            self.client.close()
            self.client = None
    def sexecute(self,command):
        feed_password = False
        transport = self.client.get_transport()
        session = transport.open_session()
        session.set_combine_stderr(False)
        #you have to check if you really need to send password here
        feed_password = self.password is not None and len(self.password) > 0
        sudocmd=command.find('sudo ')>-1
        if feed_password and sudocmd:
            session.get_pty()
        #for testing purposes we want to force sudo to always to ask for password. because of that we use "-k" key
        session.exec_command(command)
        stdin = session.makefile('wb', -1)
        stdout = session.makefile('rb', -1)
        stderr = session.makefile_stderr('rb', -1)
        
        if feed_password and sudocmd:
            stdin.write('%s\n'%self.password)
        stdin.flush()
        return {'out': stdout.readlines(),
                'err': stderr.readlines(),
                'retval': stdout.channel.recv_exit_status()}

    def execute(self, command, sudo=False):
        feed_password = False
        if sudo and self.username != "root":
            command = "sudo -S -p '' %s" % command
            feed_password = self.password is not None and len(self.password) > 0
        stdin, stdout, stderr = self.client.exec_command(command)
        if feed_password:
            stdin.write(self.password + "\n")
            stdin.flush()
        return {'out': stdout.readlines(),
                'err': stderr.readlines(),
                'retval': stdout.channel.recv_exit_status()}
    def send(self,src,dest):
        transport = paramiko.Transport((dest['host'],dest['port']))
        transport.connect(username=dest['user'],password=dest['password'])
        sftp = paramiko.SFTPClient.from_transport(transport)
        sftp.put(src,dest['file'])
        transport.close()
